package concurrent

import (
	"sync"
)

type Map[K any, V any] struct {
	*sync.Map
}

func NewMap[K any, V any]() *Map[K, V] {
	return &Map[K, V]{new(sync.Map)}
}

func (m *Map[K, V]) Load(key K) (V, bool) {
	v, ok := m.Map.Load(key)
	if ok && v != nil {
		return v.(V), true
	}
	var result V
	return result, ok
}

func (m *Map[K, V]) Store(key K, value V) {
	m.Map.Store(key, value)
}

func (m *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
	store, l := m.Map.LoadOrStore(key, value)
	if l && store != nil {
		return store.(V), true
	}
	var v V
	return v, l
}

func (m *Map[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
	v, l := m.Map.LoadAndDelete(key)
	if l && v != nil {
		return v.(V), true
	}
	var result V
	return result, l
}

func (m *Map[K, V]) Delete(key K) {
	m.Map.Delete(key)
}

func (m *Map[K, V]) Swap(key K, value V) (previous V, loaded bool) {
	p, l := m.Map.Swap(key, value)
	if l && p != nil {
		return p.(V), true
	}
	var v V
	return v, l
}

func (m *Map[K, V]) CompareAndSwap(key K, old, new V) (swapped bool) {
	return m.Map.CompareAndSwap(key, old, new)
}

func (m *Map[K, V]) CompareAndDelete(key K, old V) (deleted bool) {
	return m.Map.CompareAndDelete(key, old)
}

func (m *Map[K, V]) Range(f func(key K, value V) (shouldContinue bool)) {
	m.Map.Range(func(key, value any) bool {
		return f(key.(K), value.(V))
	})
}

func (m *Map[K, V]) Size() int {
	size := 0
	m.Map.Range(func(key, value any) bool {
		size++
		return true
	})
	return size
}
